Skip to content

perf: build spark_size LargeList lengths from i64 offsets - #5300

Open
0lai0 wants to merge 2 commits into
apache:mainfrom
0lai0:perf-5272-LargeList-offset
Open

perf: build spark_size LargeList lengths from i64 offsets#5300
0lai0 wants to merge 2 commits into
apache:mainfrom
0lai0:perf-5272-LargeList-offset

Conversation

@0lai0

@0lai0 0lai0 commented Aug 7, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #5272

Rationale for this change

Follow-up to #5233. That PR routed all list-like spark_size through Arrow's length kernel, which returns Int64 for LargeList and then needs a cast_with_options(..., Int32, safe: false) (and a to_vec() when patching null slots). That left an extra Int64 length array and Int32 cast on the LargeList path, and the bench showed it: LargeList (10% null) was only ~1.3x faster than main after #5233, while the pure List path was ~12x.

Skip the length kernel entirely for LargeList. Subtract adjacent i64 offsets straight into Int32 lengths, then apply the shared null → -1 rewrite. No intermediate Int64Array, no cast allocation.

CometSize.convert still wraps size in CASE WHEN isnotnull(child), so the production shape is the no-null path (693 ns here, within ~2x of the pure List no-null path at 392 ns).

Overflow on LargeList now surfaces as size(): list length exceeds i32::MAX (same as the scalar path) instead of the Arrow cast_with_options(safe: false) error. Spark arrays are Int-indexed and cannot exceed i32::MAX elements per row, so the overflow branch is unreachable from a Spark plan today; it is kept as a defensive guard for non-Spark producers.

What changes are included in this PR?

  • LargeList gets its own path spark_size_large_list_from_offsets, splitting off the shared spark_size_list_like (which now handles only List and FixedSizeList).
  • Hot path (when the full offset span fits in i32): offsets.windows(2).map(|w| (w[1] - w[0]) as i32).collect(). Sound because Arrow offsets are monotonically non-decreasing, so no per-row length can exceed the full span.
  • Checked fallback spark_size_large_list_lengths_checked for the rare case where the offset span exceeds i32::MAX. Uses i32::try_from per row and errors on overflow (matches the old safe: false cast contract). Skips try_from on null rows since the caller overwrites them with -1 anyway.
  • Extract the shared null → -1 rewrite into ints_with_nulls_as_neg_one so List / FixedSizeList and LargeList cannot drift.
  • Share the overflow error message via a SIZE_OVERFLOW_MSG constant used by both the array and scalar paths.
  • Bench: parameterise create_large_list_array on with_nulls, add spark_size: LargeList of long arrays and spark_size: LargeList, no nulls shapes to match the List coverage.
  • Docs: add a row to docs/source/contributor-guide/expression-audits/collection_funcs.md (date, PR, technique, speedup, benchmark file) per optimizing_expressions.md.

Benchmark (array_size, 8192 rows, on top of #5233)

shape main (post-#5233) this PR change
LargeList (10% null) 7.13 µs 1.08 µs 6.6x
LargeList of long arrays (10% null) new shape 1.06 µs —
LargeList, no nulls (production path) new shape 693 ns —

Numbers from criterion comparison against a pre-5272 baseline saved on main (macOS). List and FixedSizeList benches are unchanged as expected (their code path did not move).

How are these changes tested?

  • Existing spark_size unit tests plus three new ones:
    • test_spark_size_sliced_large_list_array: pins slicing behavior (mirrors the List slice test).
    • test_spark_size_large_list_length_overflow: single-row i32::MAX + 1 length errors, exercising the checked fallback directly via a valid OffsetBuffer.
    • test_spark_size_large_list_checked_null_row_skips_overflow: null rows with an overflowing offset delta must not error (they get rewritten to -1 by the caller, same as the fast path).
  • cargo test -p datafusion-comet-spark-expr --lib -- spark_size (14 passed).
  • cargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings.

Are there any user-facing changes?

No. Values remain bit-identical for in-range lengths; null still returns -1. The only observable change is the overflow error message text (see Rationale), which is not reachable from Spark's Int-bounded Size.

@andygrove

andygrove commented Aug 11, 2026

Copy link
Copy Markdown
Member

This review was drafted with LLM assistance (Claude Code) and edited before posting.

Nice follow-up to #5233. I verified the offset-subtraction approach against Arrow 58.4.0's own length_impl in arrow-string/src/length.rs, which uses offsets.windows(2).map(|w| w[1].sub_wrapping(w[0])).collect(). Producing i32 directly is the same shape, just skipping the Int64 return + cast. cargo test -p datafusion-comet-spark-expr --lib -- spark_size passes 14/14 on the branch.

A few things:

Missing docs change. The description mentions adding a row to docs/source/contributor-guide/expression-audits/collection_funcs.md, but only the two Rust files are in the diff. Did that get dropped in a rebase? optimizing_expressions.md asks for a dated Performance (tuned ...) line naming the technique, speedup, PR, and benchmark file. Worth noting that #5233 did not add one either, so if you're in there anyway it would be good to record both passes under ## size so the history stays complete.

Dead Int64 cast arm in spark_size_list_like. Now that LargeList has its own path, this helper only ever sees List and FixedSizeList. Arrow's length() dispatches List to length_impl::<Int32Type> and FixedSizeList builds an Int32Array directly, so lengths.data_type() can only be Int32 here. That makes the DataType::Int64 => cast_with_options(...) arm unreachable, along with the cast_with_options and CastOptions imports (not used anywhere else in the file). Could we drop the arm and the imports? You already removed the LargeList sentence from the doc comment on this function, and removing the cast machinery too would make the point of the PR visible in the code. The other => exec_err! arm still catches anything unexpected if a new list type gets routed here later.

Duplicate null-count guard in the new function.

if list.null_count() == 0 {
    return Ok(Arc::new(Int32Array::from(values)));
}
let nulls = list.nulls().unwrap();
Ok(Arc::new(ints_with_nulls_as_neg_one(values, Some(nulls))))

ints_with_nulls_as_neg_one already skips the rewrite when null_count() is zero and when nulls is None, and both branches end in Int32Array::from(values), so this collapses to Ok(Arc::new(ints_with_nulls_as_neg_one(values, list.nulls()))). Since the reason for extracting that helper was so the List and LargeList paths cannot drift, keeping a second copy of the guard here works against that.

Test coverage gaps on the branches this PR adds.

  • Both LargeList array tests use inputs with a null row, so the null_count() == 0 early return never runs under test. That is the shape you describe as the production path and benchmark at 693 ns, so it seems worth a test_spark_size_large_list_array_no_nulls mirroring the existing test_spark_size_array_no_nulls, asserting null_count() == 0 on the output.
  • For the checked fallback, test_spark_size_large_list_length_overflow pins the error and ..._checked_null_row_skips_overflow pins the null skip, but I do not see a case where a non-null row's length actually fits while the overall span overflows. That is the reason the fallback loops per row instead of just erroring outright when range > i32::MAX, so it seems worth pinning. Something like OffsetBuffer::new(vec![0i64, i32::MAX as i64, i32::MAX as i64 + 10].into()) should give [i32::MAX, 10] and would catch a future change that turns the span check into a hard error.

The overflow error text change is fine given CometSize.convert wraps the call in CASE WHEN isnotnull(child) and Spark caps arrays at Int.MaxValue. No compatibility concern there, and support levels do not change so getIncompatibleReasons() / getUnsupportedReasons() stay accurate.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Build spark_size LargeList lengths as Int32 without Int64 cast

2 participants

@0lai0 @andygrove